import json
import os, sys
import re
import random
import torch,torchaudio
if sys.platform == "darwin":
    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
now_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(now_dir)
sys.path.append(now_dir)
sys.path.append(r"D:\code\ChatTTS")
import ChatTTS
from tools.audio import pcm_arr_to_mp3_view
from tools.logger import get_logger
import openpyxl
import pandas
import soundfile as sf

logger = get_logger("Command")

#############################################
#放入chattts项目使用
# 模型下载 https://huggingface.co/2Noise/ChatTTS/tree/main/asset
#########
#############################################



class chat_reader():
    def __init__(self):
        self.chat = ChatTTS.Chat(get_logger("ChatTTS"))
        logger.info("Initializing ChatTTS...")
        if self.chat.load(source="local", compile=False, custom_path="./asset/asset"):
            logger.info("Models loaded successfully.")
        else:
            logger.error("Models load failed.")
            sys.exit(1)

    
    @classmethod
    def SplitFile(cls, file_path1, dst):
        """
        可单独使用
        将一个txt文件的内容的函数，按照第几章进行分割
        :param file_path1: txt文件路径
        :param dst: 输出路径
        :return:
        """
        # 读取音色文件，转为dict，同步记录到excel中
        if os.path.exists("./seed.xlsx"):
            pass
        else:
            print("文件不存在，请检查文件路径和文件名")
        seed_wb = pandas.read_excel("./seed.xlsx", sheet_name="Sheet1", header=0)
        seed_dict = seed_wb.to_dict(orient='records')
        # main_voice = random.choice(seed_dict)
        main_voice = "3798"  # 自定义旁白seed
        with open(file_path1, 'r', encoding="utf-8", errors="ignore") as f1:
            # 获取file的名称
            file_dir1 = file_path1.replace("\\", '/').split("/")[-1].split(".")[0]
            path1 = os.path.join(dst, file_dir1) # 保存路径
            if not os.path.exists(path1):
                os.makedirs(path1)  # 创建文件夹
            wb = openpyxl.Workbook()  # 创建一个工作簿
            current_name = os.path.join(path1, "s001.xls")   # 创建一个工作表
            sheet = wb.active
            # 获取文件每一行
            # 默认参数
            defualt_setting= {"temperature": 0.3, "top_P": 0.7, "top_K": 30, "prompt": "[speed_0]"}
            # 默认语气
            defualt_voice = '[oral_0][laugh_0][break_5]'
            for line in f1.readlines():
                try:
                    # 按照章节分割 匹配方式可以改变
                    if ("正文 【" in line) :
                        name = line.replace("正文 【","").replace("】","").strip()
                        # name = line.strip()
                        wb.save(current_name)
                        wb.close()
                        name = line.strip()     # 获取章节名
                        current_name = os.path.join(path1, name.replace(" ", "_") + ".xls")   # 每个章节重新定义一个工作表
                        wb = openpyxl.Workbook() 
                        # wb.create_sheet(name)
                        sheet = wb.active   # 获取当前活跃的工作表
                        sheet.append(["seed", "voice_param",
                                      'text_param', "content", "volume"])  # 添加标题行
                        sheet.append([main_voice,
                                      str(defualt_setting),
                                      defualt_voice, line, "10"]) 
                    else:
                        line = line.replace("\u3000", "").replace("\n", "")
                        tt1 = re.split("(“.+?”)", line)  # 按照“”分割 人物语言和旁白
                        for t in tt1:
                            t = t.strip()
                            if t != "":
                                if len(t) < 2  :
                                    print("here",t)
                                if "”" in t and "“" in t:
                                    voice_seed = str(random.choice(seed_dict)["seed"]) # 随机选择一个音色
                                    sheet.append([voice_seed,
                                                  str(defualt_setting),
                                                  defualt_voice, t, "10"])  # 人物语言
                                else:
                                    sheet.append([main_voice,
                                                  str(defualt_setting),
                                                  defualt_voice, t, "1"])  # 旁白
                except Exception as e:
                    print(e)
            wb.save(current_name) # 保存最后一个工作表
            wb.close()  # 关闭工作簿

    def save_mp3_file(self, wav, index, dir):
        """
        保存wav音频文件为MP3格式
        :param wav: 
        :param index:
        :param dir: 
        :return:
        """
        data = pcm_arr_to_mp3_view(wav)
        mp3_filename = f"output_audio_{index}.mp3"
        with open(os.path.join(dir, mp3_filename), "wb") as f:
            f.write(data)
        logger.info(f"Audio saved to {mp3_filename}")

    def speaker_seed(self, num: int):
        """
        设置 CPU 生成随机数的 种子 ，方便下次复现实验结果。
        为 CPU 设置 种子 用于生成随机数，以使得结果是确定的。
        当你设置一个随机种子时，接下来的随机算法生成数根据当前的随机种子按照一定规律生成。
        随机种子作用域是在设置时到下一次设置时。要想重复实验结果，设置同样随机种子即可。
        :param num:
        :return:
        """
        torch.manual_seed(int(num))

    def excel_read(self):
        """
        将音色配置转换为dict

        :return:
        """
        wb = pandas.read_excel("seed.xlsx", header=0)  # header=0 ,第一行作为标题
        data_dict = wb.to_dict()
        return data_dict


    def read_sentence(self, vocice_info,text,save_dir,index):
        """

        :param filename: 保存的文件名
        :param sentence_dict: {"seed":2133,"voice_param":"{}",..}
        :return:
        """
        save_path = os.path.join(save_dir, f"{index+2}.wav")


        # Perform inference and save the output audio
        print(text)

        spk = self.chat.sample_random_speaker()

        params_infer_code = ChatTTS.Chat.InferCodeParams(
        spk_emb=spk,  # add sampled speaker
        temperature=.8,  # using custom temperature
        top_P=0.9,  # top P decode
        top_K=30,  # top K decode
        prompt='[speed_3]')
        params_refine_text = ChatTTS.Chat.RefineTextParams(prompt='[oral_0][laugh_0][break_6]')

        with torch.no_grad():

            wavs = self.chat.infer(
            text,
            skip_refine_text=False,
            refine_text_only=False,  # 只refine text,不生产音频
            params_refine_text=params_refine_text,
            params_infer_code=params_infer_code)
            # wav_all = AudioSegment.empty()
            # print(len(wavs))
            # for w in wavs:
            #     wav_all += w
            torchaudio.save(save_path,torch.from_numpy(wavs[0]).unsqueeze(0),24000)

    def save_mp3_file(self, wav, mp3_filename):
        # data = pcm_arr_to_mp3_view(wav)
        with open(os.path.join(mp3_filename), "wb") as f:
            # print(len(wav))
            for w in wav:
                data = pcm_arr_to_mp3_view(w)
                data = data
                f.write(data)
        logger.info(f"Audio saved to {mp3_filename}") 


if __name__ == "__main__":
    a = chat_reader()
    # a.SplitFile(r"E:\test\混世小农民\混世小农民.txt", dst=r"E:\test\混世小农民")

    # xlxpath = r"E:\test\明朝败家子"
    # filename = r"第八百八十二章：好圣孙.xls"
    # xlsfile = os.path.join(xlxpath,filename)
    # wavname = r"e:\test\wav"

    # # 读取excel
    # seed_wb = pandas.read_excel(xlsfile, sheet_name="Sheet", header=0)
    # data_dict = seed_wb.to_dict(orient='records')

    # #TTS转换每段音频文件，保存
    # for index, sentence_dict in enumerate(data_dict):
    #     # print(data_dict["content"])
    #     name = str(index)
    #     print(index)
    #     mp3s = a.read_sentence(wavname + f"\\a{name}.mp3", sentence_dict)
    #         # for w in mp3s:
    #         #     f.write(w)
    # logger.info(f"Audio saved to {wavname}")
    
    # # 合并
    # os.chdir(wavname)
    # file_list = os.listdir(wavname)
    # file_list = natsorted(file_list)
    # print(file_list)
    # silence = AudioSegment.silent(duration=500)
    # combined_audio = silence
    # for audio in file_list:
    #     audio_sj = AudioSegment.from_file(audio, "mp3")
    #     if re.search("\_\d+",audio):
    #         num = int(re.search("\_(\d+)",audio).group(1))
    #         audio_sj = audio_sj + num
    #     combined_audio += audio_sj
    #     combined_audio += silence
    # combined_audio.export("output_audio.mp3", "mp3")
    

    vocice_info={"index":1,"prompt_speech_path":"5"}
    text= ["叫你个球！"]
    save_dir=r"e:\test\wav"
    index=1
    a.read_sentence(vocice_info=vocice_info,text=text,save_dir=save_dir,index=index)

